package org.digitarts.rtms.spring.collectors.jmx;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Iterator;
import java.util.Set;
import java.util.regex.Pattern;

import javax.management.Attribute;
import javax.management.AttributeList;
import javax.management.MBeanAttributeInfo;
import javax.management.MBeanServerConnection;
import javax.management.MBeanInfo;
import javax.management.ObjectName;
import javax.management.openmbean.CompositeData;

import org.digitarts.rtms.helpers.StringHelper;
import org.digitarts.rtms.spring.collectors.CollectorBase;


public class JMXCollector extends CollectorBase
{
	protected MBeanServerConnection			server					= null;
	protected Set<JMXCollection>			jmxCollections			= new HashSet<JMXCollection>();
	protected Map<ObjectName, MBeanInfo>	mbeanInfoCache			= new HashMap<ObjectName, MBeanInfo>();
	protected int							collectionCount			= 0;

	protected String[]						availabilityNameSpace	= null;
	protected String[]						availabilityState		= null;
	protected String[]						availabilityConnectTime	= null;

	public static Pattern					namePattern				= Pattern.compile("\\{(\\d++)\\}");

	@Override
	public void initCallback()
	{
		if (availabilityNameSpace != null)
		{
			formatSegments(availabilityNameSpace);
			availabilityState = StringHelper.append(availabilityNameSpace, "Availability");
			availabilityConnectTime = StringHelper.append(availabilityNameSpace, "Connect Time (ms)");
		}
	}

	@SuppressWarnings("unchecked")
	@Override
	public void collectCallback() throws Exception
	{
		long start = System.currentTimeMillis();
		MBeanInfo mbeanInfo = null;
		try
		{
			for (JMXCollection jmxCollection : jmxCollections)
			{
				Iterator mbeans = server.queryNames(jmxCollection.getTargetObjectName(), null).iterator();
				while (mbeans.hasNext())
				{
					try
					{
						ObjectName targetObjectName = (ObjectName) mbeans.next();
						mbeanInfo = mbeanInfoCache.get(targetObjectName);
						if (mbeanInfo == null)
						{
							mbeanInfo = server.getMBeanInfo(targetObjectName);
							mbeanInfoCache.put(targetObjectName, mbeanInfo);
						}
						AttributeList attrList = server.getAttributes(targetObjectName, jmxCollection.getAttributeNames());
						for (Object attributeObj : attrList)
						{
							Attribute attribute = (Attribute) attributeObj;
							Object originalValue = attribute.getValue();
							String name = attribute.getName();
							String type = jmxCollection.getAttributeMetricTypeMap().get(name);
							if (type == null)
							{
								throw new Exception("No metric type defined for attribute[" + name + "] for object name [" + targetObjectName + "]");
							}
							String metricName = jmxCollection.getAttributeMetricNameMap().get(name);
							if (metricName == null)
							{
								throw new Exception("No metric name defined for attribute[" + name + "] for object name [" + targetObjectName + "]");
							}
							metricName = formatName(metricName, targetObjectName);
							String[] segments = new String[jmxCollection.getSegments().length + 1];
							for (int i = 0; i < segments.length - 1; i++)
							{
								segments[i] = formatName(jmxCollection.getSegments()[i], targetObjectName);
							}
							segments[segments.length - 1] = metricName;

							if (originalValue instanceof CompositeData)
							{
								CompositeData compositeData = (CompositeData) originalValue;
								for(Object key: compositeData.getCompositeType().keySet())
								{
									tracer.trace(compositeData.get(key.toString()).toString(), type, StringHelper.append(tracingNameSpace, StringHelper.append(segments, key.toString())));
								}
							}
							else
							{
								String value = attribute.getValue().toString();
								tracer.trace(value, type, StringHelper.append(tracingNameSpace, segments));
							}

						}
						// for(String attributeName: jmxCollection.getAttributeNames()) {
						// if(this.attributeExists(attributeName, mbeanInfo)) {
						// String attrValue = server.getAttribute(targetObjectName,
						// attributeName).toString();
						// } else {
						// log.warn("The attribute [" + attributeName + "] does not exist in the
						// MBean [" + targetObjectName + "]");
						// }
						// }
					}
					catch (Exception e)
					{

					}
				}
			}
			long elapsed = System.currentTimeMillis() - start;
			tracer.trace(elapsed, elapsedTimeMetricNameSpace);
			// log.info("Collection time:" + elapsed + " ms.");
		}
		catch (Exception e)
		{
			log.error("Collect callback for [" + objectName + "] failed", e);
			throw e;
		}
	}

	protected boolean attributeExists(String attributeName, MBeanInfo mbeanInfo)
	{
		if (attributeName == null || mbeanInfo == null)
			return false;
		for (MBeanAttributeInfo attrInfo : mbeanInfo.getAttributes())
		{
			if (attrInfo.getName().equals(attributeName))
				return true;
		}
		return false;
	}

	public MBeanServerConnection getServer()
	{
		return server;
	}

	public Set<JMXCollection> getJmxCollections()
	{
		return jmxCollections;
	}

	public void setJmxCollections(Set<JMXCollection> jmxCollections)
	{
		this.jmxCollections = jmxCollections;
	}

	public int getCollectionCount()
	{
		return collectionCount;
	}

	public void setServer(MBeanServerConnection server)
	{
		this.server = server;
	}

}
